import networkx as nx
import numpy as np
from spacetrans import space
class graph:
  def __init__(self,fn):
    self.fn = fn
    self.nG = nx.DiGraph()
    self.G = nx.DiGraph()
    if fn != "":
      self.fromFile()
      self.toNum()
    else:
      pass    
  def preprocess(self):
    while( self.removeSingle()!= 0):
      pass
    while( self.removeCommon()!= 0):
      pass

  def fromFile(self):
    with open(self.fn, 'r') as f:
      for line in f:
        edge = line.strip().split()
        if len(edge) == 3:
          w = float(edge[2])
        else:
          w = 1.0
        self.G.add_edge(edge[0], edge[1], weight=w)

  def findSingle(self):
    def count(node):
      n = 0
      for nb in nx.all_neighbors(self.G,node):
        n += 1
      return n
    nn = [node for node in self.G.nodes() if count(node)==1 ]
    return nn

  def findCommon(self):
    num = nx.number_of_nodes(self.G)/20
    def count(node):
      n = 0
      for nb in nx.all_neighbors(self.G,node):
        n += 1
      return n
    nn = [node for node in self.G.nodes() if count(node)>num ]
    return nn
  def remove(self,fun):
    nn = fun()
    # print(nn)
    self.G.remove_nodes_from(nn)
    return len(nn)

  def removeSingle(self):
    return self.remove(self.findSingle)
    
  def removeCommon(self):
    return self.remove(self.findCommon)

  def toNum(self):
    k = 0
    self.name = []
    self.lookup = {}
    for n in sorted(self.G.nodes()):
      self.name.append(n)
      self.lookup[n]=k
      k += 1
    g = {}
    for s,e,d in self.G.edges(data = 'weight'):
      if self.lookup.__contains__(s) and self.lookup.__contains__(e):
        key = (self.lookup[s],self.lookup[e])
        if g.__contains__(key):
          g[key] = g[key] + 1
        else:
          g[key] = 1
    for key in g:
      self.nG.add_edge(key[0], key[1], weight=g[key])
    
  def power(self,pos):
    p = np.sum(pos[:,0]*pos[:,0])+np.sum(pos[:,1]*pos[:,1])
    pe = 0.
    for u,v,w in self.nG.edges(data='weight'):
      d = pos[u,:]-pos[v,:]
      pe += (d[0]*d[0]+d[1]*d[1])*w
    return pe/p
  
  def toNode(self):
    for u,v,d in self.nG.edges(data='weight'):
      print(u,v,d)
  def fromEMBFile(self,emb):
    sp = space()
    with open(emb, 'r') as f:
      l = 0
      for line in f:
        if l==0:
          shape = line.strip().split()
          sp.pos = np.zeros((int(shape[0]),int(shape[1])),dtype=np.float)
          sp.num = int(shape[0])
        else:
          data = line.strip().split()
          v = int(data[0])
          d = []
          for td in data[1:]:
            d.append(float(td))
          sp.pos[v] = np.array(d)
        l += 1
    return sp
  
  def indexEMBFile(self,emb):
    sp = space()
    index = []
    with open(emb, 'r') as f:
      l = 0
      for line in f:
        if l==0:
          shape = line.strip().split()
          sp.pos = np.zeros((int(shape[0]),int(shape[1])),dtype=np.float)
          sp.num = int(shape[0])
        else:
          data = line.strip().split()
          index.append(int(data[0]))
          d = []
          for td in data[1:]:
            d.append(float(td))
          sp.pos[l-1] = np.array(d)
        l += 1
    return sp,index

from sklearn.cluster import KMeans
from connect import egress,ingress

def main():
  g = graph("../test_flatten.g")
  g.toNum()
  eg = egress(g.nG)
  ig = ingress(g.nG)

  sp = g.fromEMBFile("../test_flatten128.emb")

  # sp.up(eg,ig)
  for i in range(3):
    sp.oneStep(ig)
    sp.oneStep(eg)
    print("*")
  # sp.oneStep(ig)
  # sp.oneStep(eg)
  # sp.oneStep(ig)
  # sp.oneStep(eg)
  sp.toReal()
  random_state = 170
  clusters = 20
  c_y = KMeans(n_clusters=clusters, random_state=random_state).fit_predict(sp.pos)
  # note = [[] for i in range(clusters)]
  # for (c,n) in zip(c_y,g.name):
  #   print(c,n)
  #   note[c].append(n)
  # for n in note:
  #   print(n)
  # print(c_y)
  sp.down(2)
  sp.plot(c_y)

def node(fn):
  g = graph(fn)
  g.toNum()
  g.toNode()

import sys    
if __name__ == '__main__':
  if len(sys.argv)>1:
    node(sys.argv[1])
  else:
    main()